import numpy as np
import torch
from torch.utils.data import Dataset, DataLoader
from torchvision import transforms

from Unet_train import UNet


class MedicalImageTestDataset(Dataset):
    def __init__(self, list_file):
        self.data = []
        self.label = []
        self.root_dir = list_file
        with open(list_file, 'r') as file:
            for line in file:
                image_path, label_path = line.strip().split(' ')
                self.data.append(image_path)
                self.label.append(label_path)

    def __len__(self):
        return len(self.data)

    def __getitem__(self, idx):
        img_path = self.data[idx]
        label_path = self.label[idx]
        image = np.load(img_path)
        label = np.load(label_path)
        transform = transforms.Compose([
            transforms.ToTensor(),
        ])
        image = transform(image)

        # label = transform(label).long()
        # image = image.permute(1, 0, 2)
        # label = label.permute(1, 0, 2)
        return image, label


def test(model, dataloader, device):
    model.eval()
    with torch.no_grad():
        for images, labels in dataloader:
            images = images.to(device)
            outputs = model(images)
            # 使用softmax并通过torch.argmax获取最可能的类别
            preds = torch.argmax(outputs, dim=1)
            yield images.cpu().squeeze(0).squeeze(0), preds.cpu().squeeze(0), labels.cpu().squeeze(0).squeeze(0)


if __name__ == '__main__':
    list_file = 'test_dir.txt'
    dataset = MedicalImageTestDataset(list_file)
    dataloader = DataLoader(dataset, batch_size=1, shuffle=False)  # batch_size 根据需求调整
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

    model = UNet(in_channels=1, out_channels=9)  # 确保输出通道数与训练时一致
    model.load_state_dict(torch.load('../pt_file/synapse-unet-29.pt', map_location=device, weights_only=True))
    model.to(device)

    from matplotlib.colors import Normalize
    import matplotlib.pyplot as plt

    norm = Normalize(vmin=0, vmax=8)  # 假设预测的类别标签从0到8

    for images, preds, labels in test(model, dataloader, device):
        fig, ax = plt.subplots(1, 3, figsize=(18, 6))
        images, preds, labels = images.permute(1, 0), preds.permute(1, 0), labels.permute(1, 0)
        print(torch.max(labels))
        ax[0].imshow(images, cmap='gray')
        ax[1].imshow(preds, cmap='tab20', norm=norm)
        ax[2].imshow(labels, cmap='tab20', norm=norm)

        ax[0].title.set_text('Original Image')
        ax[1].title.set_text('Predicted Mask')
        ax[2].title.set_text('True Mask')

        plt.setp(ax, xticks=[], yticks=[])
        plt.show()
